import logging
import re
import json
from typing import List, Dict, Optional, Union

# 导入文件处理库
# import docx
# import fitz as PyMuPDF

# 导入模型调用模块
import chat

# --- 配置 ---
logging.basicConfig(format="%(asctime)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S",
                    level=logging.INFO)
logger = logging.getLogger(__name__)


# --- 辅助函数 ---

def build_title_level_prompt(content: str) -> str:
    """
    构建一个指令，要求模型智能识别所有标题（包括不规范的），并返回其原文和层级。
    """
    return f"""你是一个高度智能的文档结构分析专家。你的任务是从我提供的文章内容中，识别并提取出所有的标题行。

请严格遵循以下规则：

1.  **智能识别**：分析文本的上下文、格式和内容，识别出所有作为标题的行。这至关重要，因为文本可能不遵循标准的Markdown格式。你需要识别包括但不限于：
    -   标准Markdown标题 (例如: '# 标题', '## 子标题')
    -   数字/字母/中文编号标题 (例如: '1. 标题', '1.1 子标题', '一、', '（一）', 'A.')
    -   **无标记标题**：一个单独成行的、上下文表明它是一个章节起始的短语（但请排除“摘要”、“关键词”这类非正文章节）。

2.  **判断层级**：为每个识别出的标题准确判断其逻辑层级。一级标题为1，二级为2，以此类推。

3.  **输出格式**：最终结果必须是一个严格的JSON字符串，代表一个对象数组（扁平列表）。每个对象代表一个标题，必须包含两个键：
    -   `"title"`: (字符串) **标题的完整原文**。必须**保留**原始的序号和所有空格，但请移除开头的Markdown'#'标记（如果存在）。
    -   `"level"`: (整数) 标题的层级，一级为1，二级为2，以此类推。

4.  **【至关重要的要求】**：
    -   **你的回复必须只能包含JSON文本**。
    -   **绝对不要添加任何解释、说明或代码块标记**。
    -   **你的回复必须直接以 `[` 字符开始，并以 `]` 字符结束**。
    -   示例输出: `[{{"title": "一、２０２３年农业生产与价格运行状况分析", "level": 1}}, {{"title": "（一）粮食产量再创历史新高，价格高位波动", "level": 2}}, {{"title": "１．粮食生产情况及其特点", "level": 3}}]`

现在，请处理以下文章内容：

---
{content}
---
"""


def build_tree_from_flat_list(flat_list: List[Dict]) -> List[Dict]:
    """
    核心函数：从一个包含标题和层级的扁平列表中，构建层级化的树状结构。
    """
    root = []
    path = []  # 使用一个栈来追踪当前层级的父节点

    for item in flat_list:
        if not isinstance(item, dict) or 'title' not in item or 'level' not in item:
            continue  # 跳过格式不正确的数据

        level = item['level']
        title = item['title']

        node = {
            "name": title,
            "titleNum": f"H{level}",
            "children": []
        }

        # 调整路径栈，回到正确的父节点层级
        while len(path) >= level:
            path.pop()

        # 如果路径栈为空，说明这是一个新的顶级节点
        if not path:
            root.append(node)
        # 否则，将当前节点添加到路径栈末尾节点(即父节点)的children中
        else:
            path[-1]['children'].append(node)

        # 将当前节点压入路径栈，作为后续更深层级节点的父节点
        path.append(node)

    return root


def call_chat_api_non_stream(messages: List[Dict]) -> Optional[str]:
    """以非流式方式调用chat.py中的问答API。"""
    try:
        response = chat.call_llm_api(messages, stream=False)
        if response and isinstance(response, dict):
            return response.get("choices", [{}])[0].get("message", {}).get("content", "")
        return None
    except Exception as e:
        logger.error(f"执行非流式API调用时发生错误: {e}", exc_info=True)
        return None


# def read_text_from_file(file_path: str) -> str:
#     """从 .docx 或 .pdf 文件中读取文本内容。"""
#     text = ""
#     try:
#         if file_path.lower().endswith('.docx'):
#             doc = docx.Document(file_path)
#             for para in doc.paragraphs:
#                 text += para.text + '\n'
#         elif file_path.lower().endswith('.pdf'):
#             with PyMuPDF.open(file_path) as doc:
#                 for page in doc:
#                     text += page.get_text() + '\n'
#     except Exception as e:
#         logger.error(f"从文件 {file_path} 读取文本失败: {e}", exc_info=True)
#     return text


# --- 主处理函数 ---
def process_content_to_nested_outline(content: str) -> List[Dict]:
    """
    核心处理逻辑：通过单次模型调用获取扁平列表，然后在Python中构建嵌套大纲。
    """
    if not content:
        return []

    prompt = build_title_level_prompt(content)
    messages = [{"role": "user", "content": prompt}]

    logger.info("正在调用模型提取标题和层级...")
    model_response_text = call_chat_api_non_stream(messages)

    if not model_response_text:
        logger.error("模型未能返回有效的标题列表。")
        return [{"name": "错误：模型未能返回有效内容", "titleNum": "error", "children": []}]

    try:
        json_match = re.search(r'\[.*]', model_response_text, re.DOTALL)
        if not json_match: raise ValueError("在模型返回的文本中未找到JSON数组结构")

        flat_title_list = json.loads(json_match.group(0))
        if not isinstance(flat_title_list, list): raise TypeError("解析出的JSON不是列表")

        logger.info(f"模型成功返回 {len(flat_title_list)} 个标题。正在本地构建层级结构...")

        # 调用Python函数来构建树状结构
        nested_outline = build_tree_from_flat_list(flat_title_list)

        logger.info("成功构建嵌套大纲。")
        return nested_outline

    except (json.JSONDecodeError, ValueError, TypeError) as e:
        logger.error(f"处理模型返回的JSON时失败: {e}", extra={'raw_model_output': model_response_text})
        return [{"name": f"错误：处理模型返回的JSON时失败 ({e})", "titleNum": "error", "children": []}]


def extract_outline(articles: List[Dict]) -> Union[List[Dict], Dict]:
    """文本输入主函数，处理一个或多个文本输入"""
    final_results = []
    sorted_articles = sorted(articles, key=lambda x: int(x.get('sort', 0)))

    for article in sorted_articles:
        content = article.get('content')
        title = article.get('title', '无标题')
        sort_num = str(article.get('sort', '0'))

        nested_children = process_content_to_nested_outline(content)

        article_root_node = {
            "name": title,
            "titleNum": "H1",
            "sort": sort_num,
            "children": nested_children
        }
        final_results.append(article_root_node)

    if len(articles) == 1 and final_results:
        return final_results[0].get('children', [])
    else:
        return final_results


# def extract_outline_from_file(file_path: str) -> List[Dict]:
#     """文件输入主函数，处理单个文件输入"""
    # content = read_text_from_file(file_path)
    # return process_content_to_nested_outline(content)
